from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, ForeignKey
from sqlalchemy.orm import declarative_base, sessionmaker
from datetime import datetime
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 创建简化的模型以用于数据初始化
Base = declarative_base()

class Department(Base):
    __tablename__ = "departments"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, index=True, nullable=False)
    description = Column(String(255))
    parent_id = Column(Integer)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)

class DroneType(Base):
    __tablename__ = "drone_types"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, index=True, nullable=False)
    description = Column(String(255))
    max_speed = Column(Float)
    max_payload = Column(Float)
    max_flight_time = Column(Integer)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)

class Drone(Base):
    __tablename__ = "drones"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, index=True, nullable=False)
    serial_number = Column(String(100), unique=True, index=True)
    drone_type_id = Column(Integer)
    operator_id = Column(Integer)
    status = Column(String(20), default="available")
    battery_level = Column(Integer)
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)

# 数据库连接
DATABASE_URL = "sqlite:///./pilot_manage.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 从Drones.vue中提取的无人机数据
drones_data = [
    {
        "id": 1,
        "drone_code": "DJI-001",
        "name": "无人机001",
        "drone_type_id": 1,
        "drone_type_name": "DJI Mavic 3",
        "department_id": 2,
        "department_name": "运营部",
        "status": "idle",
        "battery_level": 85,
        "created_at": "2024-01-01 00:00:00",
        "description": "主力航拍无人机"
    },
    {
        "id": 2,
        "drone_code": "DJI-002",
        "name": "无人机002",
        "drone_type_id": 2,
        "drone_type_name": "DJI Mini 3 Pro",
        "department_id": 2,
        "department_name": "运营部",
        "status": "flying",
        "battery_level": 60,
        "created_at": "2024-01-02 10:00:00",
        "description": "便携航拍无人机"
    },
    {
        "id": 3,
        "drone_code": "DJI-003",
        "name": "无人机003",
        "drone_type_id": 3,
        "drone_type_name": "DJI Air 3",
        "department_id": 1,
        "department_name": "技术部",
        "status": "maintenance",
        "battery_level": 100,
        "created_at": "2024-01-03 14:00:00",
        "description": "测试用无人机"
    },
    {
        "id": 4,
        "drone_code": "DJI-004",
        "name": "无人机004",
        "drone_type_id": 4,
        "drone_type_name": "DJI Phantom 4",
        "department_id": 3,
        "department_name": "研发部",
        "status": "fault",
        "battery_level": 20,
        "created_at": "2024-01-04 09:00:00",
        "description": "备用无人机"
    }
]

# 无人机类型数据
drone_types_data = [
    {"id": 1, "name": "DJI Mavic 3", "description": "专业航拍无人机"},
    {"id": 2, "name": "DJI Mini 3 Pro", "description": "便携式航拍无人机"},
    {"id": 3, "name": "DJI Air 3", "description": "轻巧航拍无人机"},
    {"id": 4, "name": "DJI Phantom 4", "description": "经典航拍无人机"}
]

def init_drones():
    db = SessionLocal()
    try:
        # 首先初始化无人机类型
        for drone_type_data in drone_types_data:
            existing_type = db.query(DroneType).filter(DroneType.id == drone_type_data["id"]).first()
            if not existing_type:
                drone_type = DroneType(
                    id=drone_type_data["id"],
                    name=drone_type_data["name"],
                    description=drone_type_data["description"]
                )
                db.add(drone_type)
                print(f"添加无人机类型: {drone_type_data['name']}")
            else:
                print(f"无人机类型已存在: {drone_type_data['name']}")
        db.commit()
        
        # 检查是否已存在无人机数据
        existing_drones_count = db.query(Drone).count()
        if existing_drones_count > 0:
            print(f"数据库中已存在 {existing_drones_count} 架无人机数据，跳过初始化")
            return
        
        # 初始化无人机数据
        for drone_data in drones_data:
            # 状态映射
            status_map = {
                "idle": "available",
                "flying": "busy",
                "maintenance": "maintenance",
                "fault": "broken"
            }
            
            drone = Drone(
                id=drone_data["id"],
                name=drone_data["name"],
                serial_number=drone_data["drone_code"],
                drone_type_id=drone_data["drone_type_id"],
                status=status_map.get(drone_data["status"], "available"),
                battery_level=drone_data["battery_level"],
                created_at=datetime.strptime(drone_data["created_at"], "%Y-%m-%d %H:%M:%S")
            )
            db.add(drone)
            print(f"添加无人机: {drone_data['name']}")
        
        db.commit()
        print("成功初始化 4 架无人机数据到数据库")
    except Exception as e:
        db.rollback()
        print(f"初始化无人机数据时出错: {e}")
    finally:
        db.close()

if __name__ == "__main__":
    # 创建表（如果不存在）
    Base.metadata.create_all(bind=engine)
    
    # 初始化数据
    init_drones()